package com.cll.prototype.gateway.filter;

import com.cll.prototype.gateway.dto.MethodIpPercent;
import com.cll.prototype.gateway.dto.RequestIpStrategy;
import org.apache.commons.collections4.CollectionUtils;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.cloud.gateway.filter.GatewayFilterChain;
import org.springframework.cloud.gateway.filter.GlobalFilter;
import org.springframework.core.Ordered;
import org.springframework.core.io.buffer.DataBuffer;
import org.springframework.http.HttpCookie;
import org.springframework.http.HttpStatus;
import org.springframework.http.server.PathContainer;
import org.springframework.http.server.RequestPath;
import org.springframework.http.server.reactive.ServerHttpRequest;
import org.springframework.http.server.reactive.ServerHttpResponse;
import org.springframework.stereotype.Component;
import org.springframework.web.server.ServerWebExchange;
import reactor.core.publisher.Mono;

import javax.annotation.PostConstruct;
import java.net.InetSocketAddress;
import java.util.*;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.atomic.AtomicBoolean;

/**
 * 描述信息:
 * 根据请求的IP机制，进行请求过滤
 * @author CLL
 * @version 1.0
 * @date 2020/11/9 13:42
 */
@Component
public class IpMonitorFilter implements GlobalFilter, Ordered {

    private static final Logger logger = LoggerFactory.getLogger(IpMonitorFilter.class);

    /**
     * 当前使用的限流机制信息
     * 时间窗口：访问次数
     */
    private static final List<RequestIpStrategy> STRATEGY_LIST = new ArrayList<>();

    /**
     * 记录限流机制下的请求次数
     * KEY:限流机制RequestIpStrategy的getCurrentKey
     * VALUE:限流机制下对应的时间窗口的请求执行信息
     */
    private static final Map<String, MethodIpPercent> METHOD_IP_PERCENT_MAP = new ConcurrentHashMap<>();

    /**
     * 进行限流的API集合
     */
    private static final Set<String> RATE_LIMITER_API = new HashSet<>();

    /**
     * 需要登录token才能访问的API集合
     */
    private static final Set<String> TOKEN_AUTH_API = new HashSet<>();

    /**
     * 初始化限流机制、
     * 初始化限流机制的请求信息
     */
    @PostConstruct
    public void init(){
        logger.info("===>>> 初始化限流机制和限流机制对应的时间窗口的请求信息");
        // 当前限流机制规定10s内只允许访问2次
        STRATEGY_LIST.add(new RequestIpStrategy(10, 2));

        // 根据限流机制生成限流的时间窗口
        STRATEGY_LIST.forEach((requestIpStrategy) -> METHOD_IP_PERCENT_MAP.put(requestIpStrategy.getCurrentKey(), MethodIpPercent.getInstance(requestIpStrategy.getSecond(), requestIpStrategy.getTimes())));

        // 限流的接口名称
        RATE_LIMITER_API.add("/captcha/auth_captcha/createRegisterCaptcha");
        RATE_LIMITER_API.add("/user/user/isRegistry");

        // 需要登录认证才能访问的接口
        TOKEN_AUTH_API.add("/user/login_token/getEmailByToken");
    }

    /**
     * 对请求过来的IP进行限流
     * @param exchange  请求、响应数据交换对象
     * @param chain 过滤器链
     * @return  执行结果
     */
    @Override
    public Mono<Void> filter(ServerWebExchange exchange, GatewayFilterChain chain) {
        logger.info("===>>> ip monitor filter");

        final long startTime = System.currentTimeMillis();

        ServerHttpRequest request = exchange.getRequest();
        ServerHttpResponse response = exchange.getResponse();

        InetSocketAddress remoteAddress = request.getRemoteAddress();
        if (null == remoteAddress) {
            logger.error("****** 为获取到远端请求的网卡信息，直接进行返回吧");
            response.setStatusCode(HttpStatus.FORBIDDEN);
            String responseData = "未获取到远端网卡信息，禁止访问";
            DataBuffer wrap = response.bufferFactory().wrap(responseData.getBytes());
            return response.writeWith(Mono.just(wrap));
        }
        String clientIp = remoteAddress.getHostString();

        // 获取请求的接口
        RequestPath path = request.getPath();
        PathContainer pathWithinApplication = path.pathWithinApplication();
        // 当前接口已被限流
        if (RATE_LIMITER_API.contains(pathWithinApplication.value())) {
            String methodName = generateOperateName(clientIp, pathWithinApplication.value());

            AtomicBoolean isLimited = new AtomicBoolean(false);

            METHOD_IP_PERCENT_MAP.forEach((strategyKey, methodIpPercent) -> {
                boolean limited = methodIpPercent.isLimited(methodName);
                if (limited) {
                    isLimited.set(true);
                    return;
                }
                // 如果没有被限流，记录一下当前请求
                methodIpPercent.increment(methodName, System.currentTimeMillis() - startTime);
            });
            // 当前已被限流
            if (isLimited.get()) {
                response.setStatusCode(HttpStatus.FORBIDDEN);
                String data = "访问频繁，已被限流！";
                DataBuffer wrap = response.bufferFactory().wrap(data.getBytes());
                return response.writeWith(Mono.just(wrap));
            }
        }

        // 需要登录token才能访问
        if (isAuthToken(pathWithinApplication.value())) {
            List<HttpCookie> tokenList = request.getCookies().get("token");
            if(CollectionUtils.isEmpty(tokenList)) {
                response.setStatusCode(HttpStatus.FORBIDDEN);
                String data = "未获取到token，访问已被禁止！";
                DataBuffer wrap = response.bufferFactory().wrap(data.getBytes());
                return response.writeWith(Mono.just(wrap));
            }
        }
        return chain.filter(exchange);
    }

    /**
     * 生成操作名称
     * @param clientIp  请求的IP
     * @param requestPath   请求的接口地址
     * @return  操作名称（自定义实现）
     */
    private String generateOperateName(String clientIp, String requestPath){
        return clientIp + "_" + requestPath;
    }

    private boolean isAuthToken(String requestPath){
        boolean contains = TOKEN_AUTH_API.contains(requestPath);
        if (contains) {
            return true;
        }
        AtomicBoolean result = new AtomicBoolean(false);
        TOKEN_AUTH_API.forEach((authTokenApi) -> {
            boolean startsWith = requestPath.startsWith(authTokenApi);
            if (startsWith) {
                result.set(true);
            }
        });
        return result.get();
    }

    @Override
    public int getOrder() {
        return 1;
    }
}